{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-04-02T02:08:47.750520Z",
     "start_time": "2025-04-02T02:08:43.768641Z"
    }
   },
   "source": [
    "# 感觉这块的代码可以结合着之前的nanoGPT看（Decoder-only架构）。\n",
    "\n",
    "import math, copy, time\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from modelscope.models.nlp.csanmt.translation import layer_norm\n",
    "from torch.autograd import Variable"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-02T02:55:24.833436Z",
     "start_time": "2025-04-02T02:55:24.816930Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 模型架构\n",
    "class EncoderDecoder(nn.Module):\n",
    "\n",
    "    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):\n",
    "        super(EncoderDecoder, self).__init__()      # 调用父类的nn.Module方法--Python 3 中可以简写为 super().__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.src_embed = src_embed  # 源序列（输入）的嵌入层（输入embedding+位置编码）\n",
    "        self.tgt_embed = tgt_embed  # 目标序列（输出）的嵌入层\n",
    "        self.generator = generator  # 输出生成器（通常是一个线性层 + softmax，将解码器输出映射到目标词汇表分布）\n",
    "\n",
    "    def forward(self, src, tgt, src_mask, tgt_mask):\n",
    "        \"\"\"\n",
    "        输入参数：\n",
    "            src: 源序列（输入），形状通常为 (batch_size, src_len)。\n",
    "            tgt: 目标序列（输出），形状通常为 (batch_size, tgt_len)。\n",
    "            src_mask/tgt_mask: 用于屏蔽无效位置（如填充符 <pad>）的掩码矩阵。\n",
    "        \"\"\"\n",
    "        memory = self.encode(src, src_mask)                     # 调用 encode 将 src 编码为上下文表示 memory。\n",
    "        res = self.decode(memory, src_mask, tgt, tgt_mask)      # 调用 decode 基于 memory 和 tgt 生成输出。\n",
    "        return res                      # 返回解码结果（后续由 generator 转换为概率分布）。\n",
    "\n",
    "    def encode(self, src, src_mask):\n",
    "        src_embedds = self.src_embed(src)               # 输入序列嵌入（词嵌入 + 位置编码）\n",
    "        return self.encoder(src_embedds, src_mask)      # 编码器处理\n",
    "\n",
    "    def decode(self, memory, src_mask, tgt, tgt_mask):\n",
    "        # src_mask 屏蔽输入序列中的填充符\n",
    "        # tgt_mask 防止解码器在训练时“偷看”未来词（通常结合因果掩码实现）。\n",
    "        target_embedds = self.tgt_embed(tgt)    # 将目标序列 token 转换为向量（训练时是真实目标序列，推理时是已生成的部分序列）。\n",
    "        return self.decoder(target_embedds, memory, src_mask, tgt_mask)"
   ],
   "id": "3bcef2954e59d3c5",
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-02T02:57:25.994468Z",
     "start_time": "2025-04-02T02:57:25.991161Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class Generator(nn.Module):\n",
    "\n",
    "    def __init__(self, d_model, vocab):\n",
    "        \"\"\"\n",
    "\n",
    "        :param d_model: 输入向量的维度（即 Decoder 输出的特征维度，如 512、768 等）。\n",
    "        :param vocab: 目标词汇表的大小（输出的维度，即要预测的 token 数量）。\n",
    "        \"\"\"\n",
    "        super(Generator, self).__init__()\n",
    "        self.proj = nn.Linear(d_model, vocab)   # 映射\n",
    "        # 如果 d_model=512，vocab=30000（词汇表大小），则该层的权重矩阵形状为 (512, 30000)。\n",
    "\n",
    "    def forward(self, x):\n",
    "        # 如果 x 形状是 (batch_size, seq_len, 512)，则 self.proj(x) 的形状是 (batch_size, seq_len, 30000)。\n",
    "        return F.log_softmax(self.proj(x), dim=-1)  # 对 logits 在最后一个维度（dim=-1，即词汇表维度）计算 log-softmax\n",
    "        # 输出的是 对数概率（log-probabilities），而不是直接的概率值（这样可以提高数值稳定性，并方便后续计算交叉熵损失）。"
   ],
   "id": "f9bf337eb31cd48d",
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-02T03:03:52.797357Z",
     "start_time": "2025-04-02T03:03:52.793859Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def clones(module, N):\n",
    "    \"\"\"\n",
    "     深度复制（deepcopy） 一个给定的 PyTorch 模块（module）N 次，并将这些副本存储在 nn.ModuleList 中返回\n",
    "    \"\"\"\n",
    "    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])"
   ],
   "id": "63d3edb94263a069",
   "outputs": [],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-02T05:29:02.841696Z",
     "start_time": "2025-04-02T05:29:02.835936Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class LayerNorm(nn.Module):\n",
    "    def __init__(self, feature_size, eps=1e-6):\n",
    "        \"\"\"\n",
    "        :param feature_size: 输入特征的维度（即归一化的维度大小）。\n",
    "        :param eps:\n",
    "        \"\"\"\n",
    "        super(LayerNorm, self).__init__()\n",
    "        # 如果模型认为归一化会损失信息，可以通过 a_2 放大或缩小方差，通过 b_2 调整均值。\n",
    "        self.a_2 = nn.Parameter(torch.ones(feature_size))       # 可学习的缩放参数\n",
    "        self.b_2 = nn.Parameter(torch.zeros(feature_size))      # 可学习的偏置参数\n",
    "        self.eps = eps\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        :param x: (batch_size, seq_len, feature_size) 或 (batch_size, feature_size)\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        mean = x.mean(-1, keepdim=True)     # dim=-1是对每个样本的 特征维度 归一化\n",
    "        std  = x.std(-1, keepdim=True)\n",
    "        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2"
   ],
   "id": "a408ba91f9033e10",
   "outputs": [],
   "execution_count": 20
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-02T05:37:11.881737Z",
     "start_time": "2025-04-02T05:37:11.877874Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, layer, N):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.layers = clones(layer, N)      # 深拷贝 N 个相同的 layer\n",
    "        self.norm = LayerNorm(layer.size)   # 最终的层归一化\n",
    "        \"\"\"\n",
    "        目的：\n",
    "            - 缓解深层网络的梯度消失/爆炸问题。\n",
    "            - 使输出数值分布稳定，便于后续解码器处理。\n",
    "        \"\"\"\n",
    "\n",
    "    def forward(self, x, mask):\n",
    "        \"\"\"\n",
    "        :param x:输入序列，形状通常为 (batch_size, src_seq_len, d_model)。\n",
    "        :param mask:源序列的掩码（屏蔽 <pad> 等无效位置），形状为 (batch_size, 1, src_seq_len)。\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, mask)  # 每层均传入 mask，确保自注意力机制不处理填充符（如 <pad>）。\n",
    "        return self.norm(x)"
   ],
   "id": "7cd1f3887c03135a",
   "outputs": [],
   "execution_count": 22
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-02T03:05:18.624048Z",
     "start_time": "2025-04-02T03:05:18.621456Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class SublayerConnection(nn.Module):\n",
    "    def __init__(self, size, dropout):\n",
    "        super(SublayerConnection, self).__init__()\n",
    "        self.norm = LayerNorm(size)\n",
    "        self.dropout = nn.Dropout(dropout)  # 较小的 Dropout 适合大数据集，较大的 Dropout 适合小数据集防过拟合。\n",
    "\n",
    "    def forward(self, x, sublayer):\n",
    "        \"\"\"\n",
    "        :param x: (batch_size, seq_len, d_model)\n",
    "        :param sublayer: 是一个可调用的对象（函数或者模块）\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        # 原paper方案:（Pre-LN）---> 梯度可能不稳定（尤其是深层模型），导致训练困难。\n",
    "        # x_norm = self.norm(x + self.dropout(sublayer(x)))\n",
    "        # 其他paper（Post-LN）---> 梯度更稳定，适合深层模型（如 BERT、GPT）。/ 成为现代 Transformer 的默认实现。\n",
    "        #x_norm = x + self.dropout(sublayer(self.norm(x)))  # 更加推荐\n",
    "        \"\"\"\n",
    "        为什么 Post-LN 更流行？\n",
    "        梯度流动更稳定：\n",
    "            - Pre-LN 的归一化在残差之后，深层网络中梯度容易爆炸/消失。\n",
    "            - Post-LN 的归一化在残差之前，缓解了这一问题。\n",
    "        \"\"\"\n",
    "        # 调整后的版本\n",
    "        x_norm = x + self.norm(self.dropout(sublayer(x)))\n",
    "        return x_norm\n"
   ],
   "id": "b064113793c9bf69",
   "outputs": [],
   "execution_count": 8
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-02T03:05:39.904379Z",
     "start_time": "2025-04-02T03:05:39.900479Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class EncoderLayer(nn.Module):\n",
    "    \"\"\"\n",
    "    实现了 Transformer 的 单层编码器结构，包含一个 自注意力子层 和一个 前馈神经网络子层，每个子层都通过 SublayerConnection 包装（残差连接 + 层归一化 + Dropout）\n",
    "\n",
    "    \"\"\"\n",
    "    def __init__(self, size, self_attn, feed_forward, dropout):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        self.self_attn = self_attn              # 多头自注意力模块\n",
    "        self.feed_forward = feed_forward        # 前馈神经网络模块\n",
    "        self.sublayer = clones(SublayerConnection(size, dropout), 2)    # 两个子层连接，\n",
    "        self.size = size    # model的embedding维度, 默认512\n",
    "\n",
    "    def forward(self, x, mask):\n",
    "        \"\"\"数据流向：\n",
    "        EncoderLayer = Self-Attention → Add & Norm → Feed Forward → Add & Norm\n",
    "        本代码通过 SublayerConnection 实现了 Add & Norm\n",
    "        :param x: 输入序列，形状为 (batch_size, seq_len, d_model)。\n",
    "        :param mask: 源序列的掩码（屏蔽 <pad> 等无效位置），形状为 (batch_size, 1, seq_len)。\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        # 自注意力子层  lambda x: self_attn(x, x, x, mask) 将 self_attn 转换为仅需输入 x 的函数，适配 SublayerConnection 的接口\n",
    "        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))    # self_attn(x, x, x, mask) 表示 自注意力（Q=K=V=x）。\n",
    "        # 前馈子层\n",
    "        z = self.sublayer[1](x, self.feed_forward)  # 调用 feed_forward，对自注意力的输出做非线性变换。\n",
    "        return z"
   ],
   "id": "ca9fd0042c0b603e",
   "outputs": [],
   "execution_count": 9
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-02T03:06:01.804803Z",
     "start_time": "2025-04-02T03:06:01.802035Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class Decoder(nn.Module):\n",
    "    pass"
   ],
   "id": "a62f2b5ee31f449e",
   "outputs": [],
   "execution_count": 10
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-02T03:06:20.510974Z",
     "start_time": "2025-04-02T03:06:20.507352Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class DecoderLayer(nn.Module):\n",
    "    pass"
   ],
   "id": "10afba97e02291a0",
   "outputs": [],
   "execution_count": 11
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-02T03:06:42.663008Z",
     "start_time": "2025-04-02T03:06:42.659456Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def subsequent_mask(size):\n",
    "    pass"
   ],
   "id": "1770fa78d38af5e0",
   "outputs": [],
   "execution_count": 12
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-02T03:07:21.309448Z",
     "start_time": "2025-04-02T03:07:21.305922Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def attention(query, key, value, mask=None, dropout=None):\n",
    "    pass"
   ],
   "id": "f350fbddd5acc78e",
   "outputs": [],
   "execution_count": 13
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-02T03:07:45.600653Z",
     "start_time": "2025-04-02T03:07:45.597598Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    pass"
   ],
   "id": "efaf8d43128f45af",
   "outputs": [],
   "execution_count": 14
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-02T03:08:13.828764Z",
     "start_time": "2025-04-02T03:08:13.825543Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class PositionwiseFeedForward(nn.Module):\n",
    "    pass"
   ],
   "id": "e0877112733691cc",
   "outputs": [],
   "execution_count": 15
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-02T03:08:32.214051Z",
     "start_time": "2025-04-02T03:08:32.211079Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class Embedding(nn.Module):\n",
    "    pass"
   ],
   "id": "576f3a854eab4552",
   "outputs": [],
   "execution_count": 16
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-02T03:08:59.925427Z",
     "start_time": "2025-04-02T03:08:59.923176Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "    pass"
   ],
   "id": "b426ed797b24cce9",
   "outputs": [],
   "execution_count": 17
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-04-02T03:10:00.648047Z",
     "start_time": "2025-04-02T03:10:00.643524Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 返回完整的模型\n",
    "def make_model(src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1):\n",
    "    pass"
   ],
   "id": "59a490877ce4bca4",
   "outputs": [],
   "execution_count": 18
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "",
   "id": "1fb77da73ddcc843"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
